import torch
import torch.nn as nn
from torchvision.datasets import MNIST
from torch.utils.data import DataLoader
from torchvision import models, transforms
from cddpm import DDPM, ContextUnet
from torchvision.transforms import RandomHorizontalFlip, RandomCrop, RandomRotation
import matplotlib.pyplot as plt
from torchvision.transforms.functional import to_pil_image

device = "cuda:0"
# device = "cpu"
stdict = torch.load('./data/diffusion_outputs/model_199.pth')
model = DDPM(nn_model=ContextUnet(in_channels=1, n_feat=128, n_classes=10), betas=(1e-4, 0.02), n_T=400, device="cuda:0", drop_prob=0.1)
model.load_state_dict(state_dict=stdict)
model.to(device)
tf = transforms.Compose([transforms.ToTensor()])
batch_size = 256
num_classes = 10
dataset = MNIST("./data", train=True, download=True, transform=tf)
dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False, num_workers=5)


class RDC():
    def __init__(self, model):
        '''
        + model: conditional diffusion model 
        '''
        self.ddpm = model
        self.n_T = model.n_T
        self.nn_model = model.nn_model
        self.drop_prob = model.drop_prob
        self.loss = model.loss_mse

    def expectation(self, x, c):
        '''
        + x (1, C, H, W): input image
        + c (1): label
        '''
        x = x.repeat(self.n_T, 1, 1, 1)
        c = c.repeat(self.n_T)
        noise = torch.rand_like(x)
        _ts = torch.arange(start=1, end=self.n_T + 1).to(device)
        x_t = (
            self.ddpm.sqrtab[_ts, None, None, None] * x
            + self.ddpm.sqrtmab[_ts, None, None, None] * noise
        )
        context_mask = torch.bernoulli(torch.zeros_like(c)+self.drop_prob).to(device)
        expectation = self.loss(noise, self.nn_model(x_t, c, _ts / self.n_T, context_mask))
        # print(expectation)
        return expectation.item()


def save_image(x, name):
    pil_image = to_pil_image(x.squeeze(0), mode='L')
    pil_image.save(name)


def stochastic_augment(x):
    '''
    + x (batch_size, C, H, W): input images
    '''
    # 将单通道图像转换为三通道以兼容某些变换操作
    # x = x.expand(-1, 3, -1, -1)

    augment_transforms = transforms.Compose([
        RandomHorizontalFlip(p=0.5),  # 随机水平翻转
        # transforms.RandomApply([transforms.ColorJitter(brightness=0.2)], p=0.2),  # 随机调整亮度
        transforms.RandomCrop((28, 28), padding=4),  # 随机裁剪
        RandomRotation(10, fill=(0,)),  # 随机旋转
    ])

    noisy_a = augment_transforms(x)
    noisy_b = augment_transforms(x)

    return noisy_a, noisy_b


def main():
    rdc = RDC(model)
    for images, label in dataloader:
        # a, b = stochastic_augment(images)
        # loss = model(image, label)
        preds = torch.zeros_like(label)
        for ind in range(images.shape[0]):
            image = images[ind].unsqueeze(0).to(device)
            
            logit = torch.zeros(num_classes).to(device=device)
            c = torch.arange(num_classes).to(device)
            # print(image.size(), c.size())
            for i in range(num_classes):
                logit[i] = rdc.expectation(image, c[i].unsqueeze(0))
            pred = torch.argmin(logit)
            print(ind, 'th image -- predictation: ', pred, '\tgroundtruth: ', label[ind])
            preds[ind] = pred

        correct_preds = (preds == label).sum().item()
        total_samples = preds.numel()
        accuracy = correct_preds / total_samples
        print(f"Accuracy: {accuracy * 100:.2f}%")



if __name__ == "__main__":
    main()
